from flask import Flask, request, jsonify
import requests
import json

app = Flask(__name__)

# Ollama API 的基本 URL
base_url = "http://localhost:11434/api/generate"

@app.route('/generate', methods=['POST'])
def generate_text():
    try:
        # 获取请求中的 JSON 数据
        # data = request.get_json()
        data = request.form  
        print("=====",data)   
        # 检查请求数据中是否包含必要的字段
        if 'prompt' not in data:
            return jsonify({"error": "Missing 'prompt' field in request."}), 400
        # 构建请求 Ollama 的数据
        ollama_data = {
            "model": "deepseek-r1:7b",
            "prompt": data['prompt'],
            "stream": False  # 这里暂时关闭流式响应，方便处理
        }
        # 发送请求到 Ollama
        response = requests.post(base_url, json=ollama_data)
        response.raise_for_status()
        # 解析 Ollama 的响应
        result = response.json()
        # 去掉 <think> 和 </think> 标签
        clean_response = result["response"].replace("<think>", "").replace("</think>", "")
        # 去除多余的换行符
        clean_response = clean_response.strip()
        return jsonify({"response": clean_response})
    except requests.exceptions.RequestException as e:
        return jsonify({"error": f"Request to Ollama failed: {str(e)}"}), 500
    except json.JSONDecodeError as e:
        return jsonify({"error": f"Failed to decode Ollama response: {str(e)}"}), 500
    except Exception as e:
        return jsonify({"error": f"An unexpected error occurred: {str(e)}"}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=3000)